import json
import os
import re
import ast
from openai import OpenAI
from typing import Set, List, Dict, Any


class FunctionCodeGenerator:
    """
    函数代码生成器 - 生成普通Python函数而不是MCP工具
    专为新的统一MCP架构设计
    """
    
    def __init__(self, config):
        self.client = OpenAI(
            api_key=config.get('api_key', os.getenv("DASHSCOPE_API_KEY")),
            base_url=config.get('base_url', config.get('api_base', "https://dashscope.aliyuncs.com/compatible-mode/v1"))
        )
        self.model = config.get('model_name', config.get('model', 'deepseek-v3'))
        
        # 调试信息：显示加载的配置
        print(f"FunctionCodeGenerator initialized with:")
        print(f"  Model: {self.model}")
        print(f"  Base URL: {config.get('base_url', config.get('api_base', 'default_url'))}")
        print(f"  Config keys: {list(config.keys())}")

    def generate_function_code(self, indicator_name: str, logic_details: dict) -> str:
        """
        生成纯函数代码（不包含MCP装饰器）
        
        Args:
            indicator_name: 指标名称
            logic_details: 逻辑详情
            
        Returns:
            str: 生成的函数代码
        """
        prompt = self._build_function_prompt(indicator_name, logic_details)
        
        try:
            print(f"Calling OpenAI API for indicator: {indicator_name}")
            print(f"Model: {self.model}")
            print(f"Prompt length: {len(prompt)}")
            
            completion = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": "你是一位专业的Python工程师，专门生成业务逻辑函数。\n\n**最重要的要求：确保生成的代码语法完全正确，特别注意：**\n1. 所有括号必须正确匹配 ()', [], {}\n2. 列表推导式语法必须完整\n3. f-string中的大括号必须正确转义\n4. 字符串引号必须正确匹配\n5. 缩进必须一致\n6. **类型转换安全：绝对不要直接使用 float(None)，必须先检查None值**\n\n**其他要求：**\n1. 生成纯Python函数，不包含任何装饰器\n2. 使用中文编写所有注释和文档字符串\n3. 正确理解业务逻辑中的路由分发机制\n4. 熟练使用MySQLUtil进行数据库操作\n5. 实现完整的错误处理和日志记录\n6. 确保生成的函数可以直接调用\n7. 包含所有必要的导入语句\n8. 函数参数必须有详细的类型注解和描述\n9. **数据库查询结果可能包含None值，所有类型转换前必须检查**"},
                    {"role": "user", "content": prompt}
                ]
            )
            
            print(f"OpenAI API call successful for indicator: {indicator_name}")
            raw_content = completion.choices[0].message.content
            print(f"Raw content length: {len(raw_content) if raw_content else 'None'}")
            
            # 提取代码
            code = self._extract_python_code(raw_content)
            if not code:
                print(f"ERROR: Failed to extract code for indicator: {indicator_name}")
                print(f"ERROR: Raw content preview: {raw_content[:500] if raw_content else 'None'}...")
                return ""

            print(f"Code extracted successfully for indicator: {indicator_name}, length: {len(code)}")
            
            # 语法校验
            if not self._is_valid_syntax(code):
                print(f"ERROR: Generated code for indicator '{indicator_name}' has syntax errors.")
                print(f"ERROR: Code preview: {code[:500]}...")
                return ""
                
            print(f"Code validation successful for indicator: {indicator_name}")
            return code
            
        except Exception as e:
            print(f"ERROR: Exception during code generation for indicator '{indicator_name}': {str(e)}")
            print(f"ERROR: Exception type: {type(e).__name__}")
            import traceback
            print(f"ERROR: Full traceback: {traceback.format_exc()}")
            return ""

    def _build_function_prompt(self, indicator_name: str, logic_details: dict) -> str:
        """构建函数生成的prompt"""
        func_name = "get_" + re.sub(r'\W|^(?=\d)', '_', indicator_name).lower()

        logic_description = ""
        if '项目' in logic_details:
            for project_id, logic in logic_details['项目'].items():
                logic_description += f'  - 项目编码 "{project_id}" 的计算逻辑: {logic}\n'
        if '区域' in logic_details:
            for region_project_key, logic in logic_details['区域'].items():
                logic_description += f'  - 区域项目组合 "{region_project_key}" 的计算逻辑: {logic}\n'
        if '产品' in logic_details:
            for product_id, logic in logic_details['产品'].items():
                logic_description += f'  - 产品编码 "{product_id}" 的计算逻辑: {logic}\n'
        if '通用' in logic_details:
            logic_description += f'  - 通用计算逻辑: {logic_details["通用"]}\n'

        # 优化后的 prompt 模板
        prompt_template = f'''你是一位资深Python数据工程师，任务是根据业务需求，生成一个高质量、生产级别的Python函数。

### **任务：生成指标计算函数**

**指标名称**: `{indicator_name}`

**核心计算逻辑**:
{logic_description}

### **⚠️ 关键要求：完整参数识别**
在编写函数前，必须仔细分析"核心计算逻辑"，识别所有需要的输入参数：

#### 参数识别规则：
1. **往届数据对比**：如果逻辑中提到"2022届"、"2023届"、"往届"、"历史数据"、"年度对比"等关键词
   - 必须添加：`previous_year_1_project_id: Optional[int] = None`
   - 必须添加：`previous_year_1_questionnaire_ids: Optional[List[int]] = None`
   - 如果有多个往届，继续添加 `previous_year_2_*` 等参数
   
2. **数据筛选维度**：如果逻辑中提到"各学历"、"按学历"、"不同专业"、"分院系"、"按性别"等关键词
   - "学历" → `education: Optional[str] = None`
   - "专业" → `major: Optional[str] = None`
   - "院系" → `department: Optional[str] = None`
   - "性别" → `gender: Optional[str] = None`
   
3. **分组统计**：如果逻辑中提到"分组"、"分别统计"、"按XX分类"
   - 必须添加分组字段参数：`group_by_field: Optional[str] = None`
   
4. **全国数据对比**：如果逻辑中提到"全国数据"、"与全国对比"
   - 通常不需要额外参数，在函数内查询全国数据表即可

#### 参数识别示例：
**示例1：** "计算当前项目(2024届)及往届项目（2022届和2023届）的境内升学比例"
- 分析：提到了多个届次的数据对比
- 需要参数：
  ```python
  previous_year_1_project_id: Optional[int] = None  # 2023届
  previous_year_1_questionnaire_ids: Optional[List[int]] = None
  previous_year_2_project_id: Optional[int] = None  # 2022届
  previous_year_2_questionnaire_ids: Optional[List[int]] = None
  ```

**示例2：** "统计各学历层次（本科、硕士、博士）的就业满意度"
- 分析：需要按学历筛选数据
- 需要参数：
  ```python
  education: Optional[str] = None  # 可选的学历筛选
  ```

**示例3：** "查询XX项目的总体就业率"
- 分析：只需要基础参数，无额外筛选
- 需要参数：只需核心参数（project_id, questionnaire_ids等）

#### 其他注意事项：
- **参数命名**：使用英文命名，遵循Python命名规范（小写+下划线），类型注解必须准确
- **参数默认值**：所有额外参数必须设置为 `None` 作为默认值，使用 `Optional[]` 类型注解
- **参数位置**：额外参数应放在核心参数（product_code, project_code, region_code）之后

**⚠️ 不要遗漏任何必要的参数！** 函数必须能够完整地执行"核心计算逻辑"中描述的所有功能。

---

### **代码生成要求**

#### **1. 函数签名与参数**
- **函数名**: 必须是 `{func_name}`。
- **核心参数**: 函数必须包含以下参数，并正确添加类型注解：
    - `project_id: int`: 项目ID。
    - `questionnaire_ids: List[int]`: 问卷ID列表。
    - `product_code: Optional[str] = None`: 产品编码。
    - `project_code: Optional[str] = None`: 项目编码。
    - `region_code: Optional[str] = None`: 区域编码。
- **动态参数（重要）**: 仔细分析"核心计算逻辑"，识别所有需要的额外参数。常见场景包括：
    - **数据筛选维度**: 
        - `education: Optional[str] = None` - 学历层次（本科/硕士/博士）
        - `department: Optional[str] = None` - 院系
        - `major: Optional[str] = None` - 专业
        - `gender: Optional[str] = None` - 性别
    - **往届数据对比**: 
        - `previous_year_1_project_id: Optional[int] = None` - 上一届项目ID（如2023届）
        - `previous_year_1_questionnaire_ids: Optional[List[int]] = None` - 上一届问卷ID集合
        - `previous_year_2_project_id: Optional[int] = None` - 上两届项目ID（如2022届）
        - `previous_year_2_questionnaire_ids: Optional[List[int]] = None` - 上两届问卷ID集合
    - **分组统计**: 
        - `group_by_field: Optional[str] = None` - 分组字段名称
        - `filter_condition: Optional[str] = None` - 额外筛选条件
    - **数据范围**: 
        - `start_date: Optional[str] = None` - 开始日期
        - `end_date: Optional[str] = None` - 结束日期
        - `region_filter: Optional[List[str]] = None` - 区域过滤列表
- **硬编码参数**: 计算逻辑中提到的固定值（如问题编码 'T00000373'、固定的SQL条件）必须硬编码在函数内部，不能作为参数。

#### **2. 函数文档字符串 (Docstring) -【重点】**
**必须**按照以下格式和要求，用中文编写详细的Docstring：

```python
"""
{indicator_name} - 指标计算函数

## 指标说明
[在这里详细说明该指标的业务含义、计算目的和数据内容。根据“核心计算逻辑”进行总结。]

## Args
    project_id (int): 当前项目ID，用于查询项目配置信息。
    questionnaire_ids (List[int]): 当前项目问卷ID集合，用于确定数据范围。
    product_code (Optional[str]): 产品编码，用于路由到特定计算逻辑。
    project_code (Optional[str]): 项目编码，用于路由到特定计算逻辑。
    region_code (Optional[str]): 区域编码，用于路由到特定计算逻辑。
    [根据"核心计算逻辑"，在此处补充所有必要的参数说明。参考动态参数列表。]

## 示例
### 输入示例1：基础查询（无额外参数）
```json
{{
    "project_id": 5895,
    "questionnaire_ids": [11158, 11159]
}}
```

### 输入示例2：按学历筛选
```json
{{
    "project_id": 5895,
    "questionnaire_ids": [11158, 11159],
    "education": "本科"
}}
```

### 输入示例3：往届数据对比
```json
{{
    "project_id": 6768,
    "questionnaire_ids": [12709, 12710],
    "previous_year_1_project_id": 0,
    "previous_year_1_questionnaire_ids": [0],
    "previous_year_2_project_id": 461,
    "previous_year_2_questionnaire_ids": [831, 832]
}}
```

### 输出
```json
[根据指标逻辑和输入示例，在此处生成一个真实、合理的JSON输出示例。确保输出结构与返回格式一致。]
```
"""
```

#### **3. 函数实现核心要点**
- **纯函数**: 只返回纯Python函数代码，不包含任何装饰器。
- **完整导入**: 必须包含所有必要的库导入，尤其是 `sys`, `os`, `logging`, `json`, `re`, `typing` 和 `MySQLUtil`。
- **数据库操作 (`MySQLUtil`)**:
    - **导入方式**: 必须使用以下代码导入 `MySQLUtil`:
      ```python
      import sys
      import os
      sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..'))
      from shared.utils.MySQLUtil import MySQLUtil
      ```
    - **使用**: 实例化 `db = MySQLUtil()`，使用 `db.fetchone()` 和 `db.fetchall()`。
    - **空值检查**: 所有数据库查询结果必须进行严格的空值检查（`if not result:`），避免 `NoneType` 错误。
- **动态SQL**:
    - **IN 子句**: 必须使用 f-string 和 `','.join(['%s'] * len(ids))` 的方式构造 `IN` 子句，以防止SQL注入和格式化错误。
    - **动态表名**: 答案表等表名需要从 `client_item` 表查询 `split_tb_paper` 和 `item_year` 字段后动态拼接，如 `f"re_dy_paper_answer_{{split_tb_paper}}"`。
- **错误处理**:
    - **`try...except`**: 整个函数核心逻辑必须包裹在 `try...except Exception as e:` 块中。
    - **日志记录**: 在 `except` 块中，使用 `logger.error(..., exc_info=True)` 记录完整的异常信息。在关键步骤（如查询数据库、获取配置）使用 `logger.info()` 记录过程。
- **返回结构**:
    - **成功**: `{{ "success": True, "message": "ok", "code": 0, "result": [...] }}`
    - **失败**: `{{ "success": False, "message": "错误信息", "code": 500, "error": str(e) }}`

#### **4. 常见陷阱（务必规避）**
- **参数缺失**: 这是最严重的错误！如果计算逻辑中提到多个项目、多个届次、多个学历层次等，必须添加相应的参数来接收这些信息。不要假设这些值可以从其他地方获取。
- **往届数据对比陷阱**: 如果需要对比多个届次的数据，每个届次都需要独立的 `project_id` 和 `questionnaire_ids` 参数。不要试图在函数内部硬编码项目ID。
- **问卷过滤**: 务必先根据 `dy_target` 字段过滤问卷ID，再用过滤后的ID查询问题和答案。
- **答案表字段名称**: 答案表 `re_dy_paper_answer_*` 中**不存在** `ans` 字段，必须使用 `c1, c2, c3, c4, c5` 等字段。正确写法是 `SUM(t1.c1)` 而不是 `SUM(CASE WHEN t1.ans = '1' THEN 1 ELSE 0 END)`。
- **数据类型混合运算**: 数据库返回的是 `Decimal` 类型，在数学运算时必须确保所有操作数类型一致。正确做法：
  - **错误示例**: `float(result['c1']) / total` (当 total 是 Decimal 时会报错)
  - **正确示例**: `float(result['c1'] or 0) / float(total)` 或在计算 total 时就统一转换 `sum([float(result[f'c{{i}}'] or 0) for i in range(1, 6)])`
- **空值运算**: 从字典或数据库结果取值时，使用 `result['key'] or 0` 的方式处理 `None`，防止 `None` 参与运算。同时在类型转换时也要处理空值：`float(result['c1'] or 0)`。
- **float(None)类型转换错误**: 当SQL聚合函数返回NULL时（如某学历层次无数据时SUM()/COUNT()会返回NULL），直接转换会报错。
  - **错误示例**: `{{k: float(v) for k, v in result.items()}}` (当v为None时报错: float() argument must be a string or a real number, not 'NoneType')
  - **正确示例**: `{{k: float(v) if v is not None else 0.0 for k, v in result.items()}}`
  - **SQL预防**: 使用CASE WHEN语句避免NULL值：`CASE WHEN COUNT(*) = 0 THEN 0 ELSE SUM(c1)/COUNT(*) END`
- **多选题 vs 单选题**: 多选题的答案在 `c1, c2, ...` 等多个字段，单选题在 `ans` 或 `c1` 字段，注意SQL写法的区别。

#### **5. 代码结构模板**
请严格按照以下模板结构生成代码：

```python
import sys
import os
import logging
import re
import json
from typing import Dict, List, Any, Optional, Tuple

# 添加项目根路径到 sys.path
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..'))

# 导入 MySQLUtil
from shared.utils.MySQLUtil import MySQLUtil

# 设置日志
logger = logging.getLogger(__name__)

def {func_name}(
    project_id: int,
    questionnaire_ids: List[int],
    product_code: Optional[str] = None,
    project_code: Optional[str] = None,
    region_code: Optional[str] = None,
    # 根据"核心计算逻辑"添加其他必要参数，例如：
    # education: Optional[str] = None,  # 学历筛选
    # previous_year_1_project_id: Optional[int] = None,  # 往届对比
    # previous_year_1_questionnaire_ids: Optional[List[int]] = None,
    # previous_year_2_project_id: Optional[int] = None,
    # previous_year_2_questionnaire_ids: Optional[List[int]] = None
) -> Dict[str, Any]:
    """
    {indicator_name} - 指标计算函数
    
    ## 指标说明
    [详细说明该指标的业务含义、计算目的和数据内容]
    
    ## Args
        project_id (int): 当前项目ID，用于查询项目配置信息
        questionnaire_ids (List[int]): 当前项目问卷ID集合，用于确定数据范围
        product_code (Optional[str]): 产品编码，用于路由到特定计算逻辑
        project_code (Optional[str]): 项目编码，用于路由到特定计算逻辑
        region_code (Optional[str]): 区域编码，用于路由到特定计算逻辑
        [如果有额外参数，在此处补充说明]
        
    ## 示例
    ### 输入
    ```json
    {{
        "project_id": 5895,
        "questionnaire_ids": [11158, 11159]
    }}
    ```
    
    ### 输出
    ```json
    {{
        "success": true,
        "message": "ok", 
        "code": 0,
        "result": [...]
    }}
    ```
    """
    logger.info(f"开始计算指标: {indicator_name}, 项目ID: {{project_id}}")
    
    try:
        db = MySQLUtil()  

        # 1. 查询项目配置信息
        project_sql = """
        SELECT client_code, item_year, dy_target_items, split_tb_paper 
        FROM client_item 
        WHERE id = %s
        """
        project_info = db.fetchone(project_sql, (project_id,))
        if not project_info:
            raise ValueError(f"未找到项目ID={{project_id}}的配置信息")

        client_code = project_info['client_code']
        item_year = project_info['item_year']
        split_tb_paper = project_info['split_tb_paper']
        
        logger.info(f"项目配置: client_code={{client_code}}, item_year={{item_year}}, split_tb_paper={{split_tb_paper}}")

        # 2. 计算 shard_tb_key
        shard_tb_key = re.sub(r'^[A-Za-z]*0*', '', client_code)
        logger.info(f"计算得到 shard_tb_key: {{shard_tb_key}}")

        # 3. 查询问卷信息
        questionnaire_sql = f"""
        SELECT id, dy_target 
        FROM wt_template_customer 
        WHERE id IN ({{','.join(['%s'] * len(questionnaire_ids))}})
        """
        questionnaires = db.fetchall(questionnaire_sql, tuple(questionnaire_ids))
        if not questionnaires:
            raise ValueError(f"未找到问卷ID集合={{questionnaire_ids}}的配置信息")
        
        logger.info(f"查询到问卷信息: {{questionnaires}}")

        # 4. 过滤特定调研对象的问卷（根据实际需要修改 dy_target）
        valid_questionnaire_ids = [q['id'] for q in questionnaires if q['dy_target'] == 'GRADUATE_SHORT']
        if not valid_questionnaire_ids:
            raise ValueError("未找到目标调研对象的问卷ID")
            
        logger.info(f"找到有效问卷ID: {{valid_questionnaire_ids}}")

        # 5. 查询问题信息（根据实际问题编码修改）
        question_sql = """
        SELECT id, wt_code, wt_obj 
        FROM wt_template_question_customer 
        WHERE cd_template_id = %s AND wt_code = %s AND is_del = 0
        """
        question_info = db.fetchone(question_sql, (valid_questionnaire_ids[0], 'T00000373'))
        if not question_info:
            raise ValueError("未找到指定问题编码的问题信息")
            
        logger.info(f"找到问题信息: {{question_info['id']}}")

        # 6. 解析问题选项（如果需要）
        wt_obj = json.loads(question_info['wt_obj'])
        options = []
        for item in wt_obj['itemList']:
            options.append({{
                'key': item['key'],
                'val': item['val'],
                'weight': item.get('weight', 1)
            }})

        # 7. 构建动态表名
        answer_table = f"re_dy_paper_answer_{{split_tb_paper}}"
        student_table = f"dim_client_target_baseinfo_student_{{item_year}}"

        # 8. 路由分发逻辑
        def _calc_default():
            """默认计算逻辑"""
            # 根据具体指标实现计算逻辑，以下只是示例
            # 注意：答案表中使用 c1, c2, c3, c4, c5 字段，不是 ans 字段
            sql = f"""
            SELECT 
                SUM(t1.c1) as c1,
                SUM(t1.c2) as c2,
                SUM(t1.c3) as c3,
                SUM(t1.c4) as c4,
                SUM(t1.c5) as c5
            FROM {{answer_table}} t1
            WHERE t1.cd_template_id = %s AND t1.wid = %s AND t1.ans_true = 1
            """
            result = db.fetchone(sql, (valid_questionnaire_ids[0], question_info['id']))
            if not result:
                raise ValueError("未找到有效的答案数据")
            
            # 计算总数和比例（示例：前三项占比）
            # 注意：使用 or 0 处理None值，避免float(None)错误
            total = sum([float(result[f'c{{i}}'] or 0) for i in range(1, 6)])
            if total == 0:
                ratio = 0
            else:
                ratio = float((result['c1'] or 0) + (result['c2'] or 0) + (result['c3'] or 0)) / total
            
            # 处理查询结果，构建返回数据
            return [{{
                "key": "示例比例",
                "val": f"{{ratio:.6f}}"
            }}]

        # 路由映射
        project_logic_map = {{
            # "项目编码": lambda: _calc_for_project_xxx()
        }}
        
        # 按优先级路由
        result_data = None
        if project_code and project_code in project_logic_map:
            result_data = project_logic_map[project_code]()
        else:
            result_data = _calc_default()

        logger.info(f"指标 '{indicator_name}' 计算成功")
        return {{
            "success": True,
            "message": "ok",
            "code": 0,
            "result": [{{
                "calcVals": [{{
                    "val": result_data
                }}]
            }}]
        }}

    except Exception as e:
        logger.error(f"计算指标 '{indicator_name}' 时发生错误: {{str(e)}}", exc_info=True)
        return {{
            "success": False,
            "message": f"数据获取失败: {indicator_name}",
            "code": 500,
            "error": str(e)
        }}
```

---

### **最终交付**

你的回复必须是**一段完整的Python代码**，包含所有导入、函数定义和文档，并用单个 ```python ... ``` 代码块包裹。**严禁**在代码块之外添加任何解释性文字。
'''
        return prompt_template

    def _extract_python_code(self, content: str) -> str:
        """提取Python代码块"""
        # 首先尝试标准的markdown代码块格式
        match = re.search(r"```python\n(.*?)\n```", content, re.DOTALL)
        if match:
            code = match.group(1).strip()
        else:
            # 如果没有找到标准格式，使用整个内容
            code = content.strip()
        
        # 清理可能的错误前缀
        lines = code.split('\n')
        cleaned_lines = []
        
        for line in lines:
            # 跳过只包含 'python' 的行
            if line.strip() == 'python':
                continue
            # 跳过空行（在开头）
            if not cleaned_lines and not line.strip():
                continue
            cleaned_lines.append(line)
        
        return '\n'.join(cleaned_lines)

    def _is_valid_syntax(self, code: str) -> bool:
        """验证代码语法"""
        try:
            ast.parse(code)
            return True
        except SyntaxError as e:
            print(f"ERROR: Syntax validation failed: {e}")
            print(f"ERROR: Syntax error line: {e.lineno}")
            print(f"ERROR: Syntax error text: {e.text}")
            print(f"ERROR: Syntax error offset: {e.offset}")
            # 显示出错前后的代码行
            lines = code.split('\n')
            start_line = max(0, e.lineno - 3) if e.lineno else 0
            end_line = min(len(lines), e.lineno + 2) if e.lineno else 5
            print(f"ERROR: Code context around error:")
            for i in range(start_line, end_line):
                marker = " >>> " if i == e.lineno - 1 else "     "
                print(f"{marker}{i+1:3d}: {lines[i] if i < len(lines) else ''}")
            return False
        except Exception as e:
            print(f"ERROR: Unexpected error during syntax validation: {e}")
            print(f"ERROR: Error type: {type(e).__name__}")
            return False

    
